Skip to content

perf: batch metadata query in KB retrieval to fix N+1 problem#5463

Merged
Soulter merged 2 commits intoAstrBotDevs:masterfrom
CAICAIIs:master
Feb 26, 2026
Merged

perf: batch metadata query in KB retrieval to fix N+1 problem#5463
Soulter merged 2 commits intoAstrBotDevs:masterfrom
CAICAIIs:master

Conversation

@CAICAIIs
Copy link
Contributor

@CAICAIIs CAICAIIs commented Feb 25, 2026

Modifications / 改动点

  • astrbot/core/knowledge_base/kb_db_sqlite.py:新增 get_documents_with_metadata_batch(doc_ids: set[str]) 方法,使用 SQL IN 子句实现批量文档元数据查询,返回 {doc_id: metadata_dict} 映射
  • astrbot/core/knowledge_base/retrieval/manager.py:修改 RetrievalManager.retrieve() 方法,将 N+1 循环查询替换为单次批量查询调用
  • This is NOT a breaking change. / 这不是一个破坏性变更。

Screenshots or Test Results / 运行截图或测试结果

Benchmark results (local SQLite):

  • 10 docs: 10.67ms → 1.47ms (7.3x faster)
  • 20 docs: 26.00ms → 2.68ms (9.7x faster)
  • 50 docs: 63.87ms → 2.79ms (22.9x faster)

Checklist / 检查清单

  • 😊 如果 PR 中有新加入的功能,已经通过 Issue / 邮件等方式和作者讨论过。/ If there are new features added in the PR, I have discussed it with the authors through issues/emails, etc.
  • 👀 我的更改经过了良好的测试,并已在上方提供了“验证步骤”和“运行截图”。/ My changes have been well-tested, and "Verification Steps" and "Screenshots" have been provided above.
  • 🤓 我确保没有引入新依赖库,或者引入了新依赖库的同时将其添加到了 requirements.txtpyproject.toml 文件相应位置。/ I have ensured that no new dependencies are introduced, OR if new dependencies are introduced, they have been added to the appropriate locations in requirements.txt and pyproject.toml.
  • 😮 我的更改没有引入恶意代码。/ My changes do not introduce malicious code.

Summary by Sourcery

通过批量查询文档元数据来优化知识库检索,从而消除 N+1 访问模式。

增强内容:

  • 为 SQLite 知识库存储添加批量文档及知识库元数据检索 API。
  • 更新检索管理器,在融合结果时使用批量元数据查询替代按文档逐个查询。
Original summary in English

Summary by Sourcery

Optimize knowledge base retrieval by batching document metadata queries to eliminate an N+1 access pattern.

Enhancements:

  • Add a batch document-and-knowledge-base metadata retrieval API to the SQLite KB store.
  • Update retrieval manager to use batched metadata lookup instead of per-document queries for fused results.

Replace N sequential get_document_with_metadata() calls with a single
get_documents_with_metadata_batch() call using SQL IN clause.

Benchmark results (local SQLite):
- 10 docs: 10.67ms → 1.47ms (7.3x faster)
- 20 docs: 26.00ms → 2.68ms (9.7x faster)
- 50 docs: 63.87ms → 2.79ms (22.9x faster)
@auto-assign auto-assign bot requested review from Fridemn and Raven95676 February 25, 2026 13:21
@dosubot dosubot bot added the size:M This PR changes 30-99 lines, ignoring generated files. label Feb 25, 2026
@gemini-code-assist
Copy link
Contributor

Summary of Changes

Hello @CAICAIIs, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

此拉取请求旨在通过优化知识库文档元数据的查询机制来大幅提升系统性能。通过引入批量查询方法并将其集成到文档检索流程中,成功解决了 N+1 查询问题,从而显著减少了数据库交互次数和响应时间。

Highlights

  • 批量元数据查询: 在 astrbot/core/knowledge_base/kb_db_sqlite.py 中新增了 get_documents_with_metadata_batch 方法,该方法利用 SQL IN 子句实现文档元数据的批量查询,显著提升了查询效率。
  • 解决 N+1 问题: 修改了 astrbot/core/knowledge_base/retrieval/manager.py 中的 RetrievalManager.retrieve 方法,将原有的 N+1 循环查询替换为对新批量查询方法的单次调用,从而消除了性能瓶颈。
  • 性能优化: 经过基准测试,此优化在不同文档数量下带来了显著的性能提升,例如处理 50 个文档时速度提升了 22.9 倍。
Changelog
  • astrbot/core/knowledge_base/kb_db_sqlite.py
    • 新增了 get_documents_with_metadata_batch 异步方法,用于根据文档 ID 列表批量获取文档及其所属知识库的元数据。
    • 该方法使用 SQLAlchemy 的 select 和 join 操作,并通过 where(col(KBDocument.doc_id).in_(doc_ids)) 实现批量查询。
  • astrbot/core/knowledge_base/retrieval/manager.py
    • 在 retrieve 方法中,将获取文档元数据的逻辑从逐个查询 (get_document_with_metadata) 更改为批量查询 (get_documents_with_metadata_batch)。
    • 首先收集所有需要查询元数据的 doc_id,然后一次性调用批量查询方法获取所有元数据,并通过映射 (metadata_map) 快速查找。
Activity
  • 目前没有检测到此拉取请求的评论、审查或特定进展活动。
Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature Command Description
Code Review /gemini review Performs a code review for the current pull request in its current state.
Pull Request Summary /gemini summary Provides a summary of the current pull request in its current state.
Comment @gemini-code-assist Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help /gemini help Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

  1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution.

@dosubot
Copy link

dosubot bot commented Feb 25, 2026

Related Documentation

Checked 1 published document(s) in 1 knowledge base(s). No updates required.

How did I do? Any feedback?  Join Discord

@dosubot dosubot bot added the feature:knowledge-base The bug / feature is about knowledge base label Feb 25, 2026
Copy link
Contributor

@sourcery-ai sourcery-ai bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hey —— 我发现了 1 个问题,并给出了一些总体反馈:

  • 新的 get_documents_with_metadata_batch API 可以接受更通用的 Iterable[str](或 Collection[str])而不是 list[str],这样可以避免在调用方已经有 set 或其他 ID 可迭代对象时,还被迫构造一个 list。
  • 如果在某些使用场景下 doc_ids 可能会变得很大,建议将 IN 查询拆分为多个批次,以避免在单条语句中构造非常大的 IN 子句,从而触及 SQLite 的参数数量限制或影响性能。
供 AI Agents 使用的提示
Please address the comments from this code review:

## Overall Comments
- The new `get_documents_with_metadata_batch` API could accept a more general `Iterable[str]` (or `Collection[str]`) instead of `list[str]` to avoid forcing callers to materialize a list when they already have a set or other iterable of IDs.
- If `doc_ids` can grow large in some usage patterns, consider chunking the `IN` query into batches to avoid constructing a very large `IN` clause in a single statement, which can hit SQLite parameter limits or impact performance.

## Individual Comments

### Comment 1
<location path="astrbot/core/knowledge_base/kb_db_sqlite.py" line_range="281" />
<code_context>
+                    KnowledgeBase,
+                    col(KBDocument.kb_id) == col(KnowledgeBase.kb_id),
+                )
+                .where(col(KBDocument.doc_id).in_(doc_ids))
+            )
+            result = await session.execute(stmt)
</code_context>
<issue_to_address>
**issue (performance):** Large `doc_ids` inputs may lead to very large `IN` clauses and performance issues.

If `doc_ids` can be large, this `IN` clause may both degrade query performance and exceed SQLite’s parameter limits. Consider either batching `doc_ids` into smaller chunks and unioning the results, or enforcing a maximum size for `doc_ids` before building the query.
</issue_to_address>

Sourcery 对开源项目免费 —— 如果你觉得我们的评审有帮助,请考虑分享给他人 ✨
帮我变得更有用!请在每条评论上点击 👍 或 👎,我会根据你的反馈改进后续的评审。
Original comment in English

Hey - I've found 1 issue, and left some high level feedback:

  • The new get_documents_with_metadata_batch API could accept a more general Iterable[str] (or Collection[str]) instead of list[str] to avoid forcing callers to materialize a list when they already have a set or other iterable of IDs.
  • If doc_ids can grow large in some usage patterns, consider chunking the IN query into batches to avoid constructing a very large IN clause in a single statement, which can hit SQLite parameter limits or impact performance.
Prompt for AI Agents
Please address the comments from this code review:

## Overall Comments
- The new `get_documents_with_metadata_batch` API could accept a more general `Iterable[str]` (or `Collection[str]`) instead of `list[str]` to avoid forcing callers to materialize a list when they already have a set or other iterable of IDs.
- If `doc_ids` can grow large in some usage patterns, consider chunking the `IN` query into batches to avoid constructing a very large `IN` clause in a single statement, which can hit SQLite parameter limits or impact performance.

## Individual Comments

### Comment 1
<location path="astrbot/core/knowledge_base/kb_db_sqlite.py" line_range="281" />
<code_context>
+                    KnowledgeBase,
+                    col(KBDocument.kb_id) == col(KnowledgeBase.kb_id),
+                )
+                .where(col(KBDocument.doc_id).in_(doc_ids))
+            )
+            result = await session.execute(stmt)
</code_context>
<issue_to_address>
**issue (performance):** Large `doc_ids` inputs may lead to very large `IN` clauses and performance issues.

If `doc_ids` can be large, this `IN` clause may both degrade query performance and exceed SQLite’s parameter limits. Consider either batching `doc_ids` into smaller chunks and unioning the results, or enforcing a maximum size for `doc_ids` before building the query.
</issue_to_address>

Sourcery is free for open source - if you like our reviews please consider sharing them ✨
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

此 PR 通过批量查询元数据来修复 N+1 问题,这是一个有效的性能优化。新的 get_documents_with_metadata_batch 方法在 kb_db_sqlite.py 中实现,并在 manager.py 中被调用。我的反馈主要集中在改进新方法的类型提示,以使其与调用代码更一致并略微提高效率。

Address review feedback:
- Change doc_ids param from list[str] to set[str] to avoid unnecessary conversion
- Chunk IN clause into batches of 900 to stay under SQLite's 999 parameter limit
- Remove list() wrapping at call site, pass set directly
@dosubot dosubot bot added the lgtm This PR has been approved by a maintainer label Feb 26, 2026
@Soulter Soulter merged commit 39b9e55 into AstrBotDevs:master Feb 26, 2026
6 checks passed
@astrbot-doc-agent
Copy link

No docs changes were generated in this run (docs repo had no updates).

Docs repo: AstrBotDevs/AstrBot-docs
Trigger: PR merged


AI change summary (not committed):

  • PR perf: batch metadata query in KB retrieval to fix N+1 problem #5463 为知识库检索的内部性能优化,使用批量元数据查询替代 N+1 查询
  • 不涉及用户界面、配置项或公开 API 变更
  • 知识库使用文档(zh/use/knowledge-base.md、en/use/knowledge-base.md)无需修改
  • 本次为纯内部优化,文档仓库无文件改动

Experimental bot notice:

  • This output is generated by AstrBot-Doc-Agent for review only.
  • It does not represent the final documentation form.

@ChenZihua-cn
Copy link

cccc

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

feature:knowledge-base The bug / feature is about knowledge base lgtm This PR has been approved by a maintainer size:M This PR changes 30-99 lines, ignoring generated files.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants